﻿using FaissOperation.FaissIndex.Native;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading.Tasks;

namespace FaissOperation.FaissIndex.Index
{
    public static class FaissUtils
    {
        public static float[] Flatten(float[][] vectors, int dimension)
        {
            if (vectors == null)
                throw new ArgumentNullException(nameof(vectors));

            var result = new float[vectors.Length * dimension];

            for (int i = 0; i < vectors.Length; i++)
            {
                if (vectors[i] == null)
                    throw new ArgumentNullException($"vectors[{i}] is null");

                if (vectors[i].Length != dimension)
                    throw new ArgumentException($"Vector at index {i} has length {vectors[i].Length}, expected {dimension}");

                Buffer.BlockCopy(vectors[i], 0, result, i * dimension * sizeof(float), dimension * sizeof(float));
            }

            return result;
        }

        public static (long[] ids, float[] distances) SafeSearch(IntPtr indexPtr, float[] query, int topK, int dimension)
        {
            if (indexPtr == IntPtr.Zero)
                throw new ObjectDisposedException("FaissIndex");

            if (query == null)
                throw new ArgumentNullException(nameof(query));

            if (query.Length != dimension)
                throw new ArgumentException($"Query length {query.Length} doesn't match dimension {dimension}");

            var labels = new long[topK];
            var distances = new float[topK];

            var queryHandle = GCHandle.Alloc(query, GCHandleType.Pinned);
            var distHandle = GCHandle.Alloc(distances, GCHandleType.Pinned);
            var labelHandle = GCHandle.Alloc(labels, GCHandleType.Pinned);

            try
            {
                var result = FaissNative.faiss_Index_search(indexPtr, 1, query, topK, distances, labels);
                if (result != 0)
                    throw new Exception($"faiss_Index_search failed with error code: {result}");
            }
            finally
            {
                queryHandle.Free();
                distHandle.Free();
                labelHandle.Free();
            }

            return (labels, distances);
        }

        public static void SafeAdd(IntPtr indexPtr, float[] vector, int dimension)
        {
            if (indexPtr == IntPtr.Zero)
                throw new ObjectDisposedException("FaissIndex");

            if (vector == null)
                throw new ArgumentNullException(nameof(vector));

            if (vector.Length != dimension)
                throw new ArgumentException($"Vector length {vector.Length} doesn't match dimension {dimension}");

            var handle = GCHandle.Alloc(vector, GCHandleType.Pinned);
            try
            {
                var result = FaissNative.faiss_Index_add(indexPtr, 1, vector);
                if (result != 0)
                    throw new Exception($"faiss_Index_add failed with error code: {result}");
            }
            finally
            {
                handle.Free();
            }
        }

        public static void SafeBatchAdd(IntPtr indexPtr, float[][] vectors, int dimension)
        {
            if (indexPtr == IntPtr.Zero)
                throw new ObjectDisposedException("FaissIndex");

            if (vectors == null || vectors.Length == 0)
                throw new ArgumentException("Vectors array cannot be null or empty");

            var flat = Flatten(vectors, dimension);

            var handle = GCHandle.Alloc(flat, GCHandleType.Pinned);
            try
            {
                var result = FaissNative.faiss_Index_add(indexPtr, vectors.Length, flat);
                if (result != 0)
                    throw new Exception($"faiss_Index_add failed with error code: {result}");
            }
            finally
            {
                handle.Free();
            }
        }

        public static void SafeRebuild(IntPtr indexPtr)
        {
            if (indexPtr == IntPtr.Zero)
                throw new ObjectDisposedException("FaissIndex");

            var result = FaissNative.faiss_Index_reset(indexPtr);
            if (result != 0)
                throw new Exception($"faiss_Index_reset failed with error code: {result}");
        }
    }
}
